#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@gmail.com
@time: 2018/1/22 15:04
@desc: 
"""
import logging
import time

import cv2
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import tensorflow as tf
from skimage import transform

from crowdcounting.test.tools import convert_density_map
from crowdcounting.test.tools import preprocess_image
from crowdcounting.test.tools import rotate_image


def predict_image(image, model='', sub_shape=(320, 320), input_shape=(256, 256), scale=1.0):
    """
    预测图像中的人群密度和数量

    Parameters
    ----------
    image: 待预测的图像，可以为图像路径或者numpy的array
    model: 模型的路径
    sub_shape: 大图切割的子图的大小
    input_shape: 输入网络的图像大小
    scale: float, resize_shape / input_shape, 取值(0,1]. 决定子图在缩小时相当于input_shape的大小

    Returns
    -------
    mask: density map
    """
    start_time = time.time()
    mask_factor = np.load('../../data/mask.npy')
    shape = (1280, 1280)
    resize_shape = (int(round(input_shape[0] * scale)), int(round(input_shape[1] * scale)))
    img = cv2.imread(image) if isinstance(image, str) else image
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, shape) if not (img.shape[0] == img.shape[1] == 1280) else img
    stage_time = time.time()
    logging.info('time: {:.3f}s, loading image completed !'.format(stage_time - start_time))
    start_time = stage_time

    config = tf.ConfigProto(device_count={"CPU": 4, "GPU": 1})
    config.gpu_options.allow_growth = True
    with tf.Session(config=config) as sess:
        meta_graph_def = tf.saved_model.loader.load(sess, ['frozen_model'], model)
        signature = meta_graph_def.signature_def
        image_tensor_name = signature['predict_images'].inputs['images'].name
        pre_density_map_name = signature['predict_images'].outputs['density_map'].name
        image_tensor = sess.graph.get_tensor_by_name(image_tensor_name)
        density_map_name = sess.graph.get_tensor_by_name(pre_density_map_name)
        stage_time = time.time()
        logging.info('Time: {:.3f}s, loading model completed !'.format(stage_time - start_time))
        start_time = stage_time

        with tf.device("/gpu:0"):
            mask = np.zeros(shape).astype(np.float32)
            for angle in range(0, 180, 30):
                offset, img_buf = rotate_image(angle, img)
                img1 = np.expand_dims(preprocess_image(img_buf[40:360, 480:800, :],
                                                       shape=input_shape, resize_shape=resize_shape), axis=0)
                img2 = np.expand_dims(preprocess_image(img_buf[920:1240, 480:800, :],
                                                       shape=input_shape, resize_shape=resize_shape), axis=0)
                density_map1 = sess.run([density_map_name], feed_dict={image_tensor: img1})
                density_map2 = sess.run([density_map_name], feed_dict={image_tensor: img2})
                buffer = np.zeros(shape)
                buffer[40:360, 480:800] = convert_density_map(np.squeeze(density_map1),
                                                              shape=sub_shape, resize_shape=resize_shape)
                buffer[920:1240, 480:800] = convert_density_map(np.squeeze(density_map2),
                                                                shape=sub_shape, resize_shape=resize_shape)
                if 0 != angle:
                    buffer = transform.rotate(buffer, -angle, resize=True, preserve_range=True)
                    buffer = buffer[offset[0]:(offset[0] + shape[0]), offset[1]:(offset[1] + shape[1])]
                mask += buffer
                stage_time = time.time()
                logging.info('Time: {:.3f}s, processed image of outside with degree {} !'.format(
                    stage_time - start_time, angle))
                start_time = stage_time

            small_shape = (720, 720)
            small_img = img[280:1000, 280:1000, :]
            for angle in range(0, 360, 72):
                offset, small_img_buf = rotate_image(angle, small_img)
                img_patch = np.expand_dims(preprocess_image(small_img_buf[80:400, 240:560, :],
                                                            shape=input_shape, resize_shape=resize_shape), axis=0)
                density_map = sess.run([density_map_name], feed_dict={image_tensor: img_patch})
                buffer = np.zeros(small_shape)
                buffer[80:400, 240:560] = convert_density_map(np.squeeze(density_map),
                                                              shape=sub_shape, resize_shape=resize_shape)
                if 0 != angle:
                    buffer = transform.rotate(buffer, -angle, resize=True, preserve_range=True)
                    buffer = buffer[offset[0]:(offset[0] + small_shape[0]), offset[1]:(offset[1] + small_shape[1])]
                mask[280:1000, 280:1000] += buffer
                stage_time = time.time()
                logging.info('Time: {:.3f}s, processed image of inside with degree {} !'.format(
                    stage_time - start_time, angle))

            mask *= mask_factor
            logging.info('total predict {} fisheye.'.format(np.sum(mask)))
            plt.subplot(122)
            plt.title('predict')
            mask = cv2.GaussianBlur(mask, (15, 15), 0.5)
            sns.heatmap(mask * 5, vmin=0, vmax=np.amax(mask), cbar=False)
            plt.subplot(121)
            plt.title('origin')
            plt.imshow(img)
            plt.show()
            return mask


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    predict_image('D:/workspace/data/val/testimage/535_70222017-10-27.jpg',
                  'D:/workspace/data/val/model/export/73205/', scale=0.7)
